import rosbag
import rospy
from unitree_legged_msgs.msg import LowState

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, default='')
parser.add_argument('-o', '--output', type=str, default='')

args = parser.parse_args()
input_bag = args.input
output_bag = args.output

with rosbag.Bag(output_bag, 'w') as outbag:
    for topic, msg, t in rosbag.Bag(input_bag).read_messages():
        if topic == '/ros_udp_comm_low/low_state':
            new_msg = LowState()
            new_msg.imu = msg.imu
            new_msg.motorState = msg.motorState
            new_msg.footForce = msg.footForce
            outbag.write(topic, new_msg, t)
        else:
            outbag.write(topic, msg, msg.header.stamp if msg._has_header else t)
